/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <melon/tracing/async_stack.h>

#include <unordered_set>
#include <turbo/log/logging.h>

#include <melon/portability/gmock.h>
#include <melon/portability/gtest.h>

TEST(AsyncStack, ScopedAsyncStackRoot) {
  void* const stackFramePtr = MELON_ASYNC_STACK_FRAME_POINTER();
  void* const returnAddress = MELON_ASYNC_STACK_RETURN_ADDRESS();

  KCHECK(melon::tryGetCurrentAsyncStackRoot() == nullptr);

  {
    melon::detail::ScopedAsyncStackRoot scopedRoot{
        stackFramePtr, returnAddress};
    auto* root = melon::tryGetCurrentAsyncStackRoot();
    KCHECK_NOTNULL(root);

    melon::AsyncStackFrame frame;
    scopedRoot.activateFrame(frame);

    KCHECK_EQ(root, frame.getStackRoot());
    KCHECK_EQ(stackFramePtr, root->getStackFramePointer());
    KCHECK_EQ(returnAddress, root->getReturnAddress());
    KCHECK_EQ(&frame, root->getTopFrame());

    melon::deactivateAsyncStackFrame(frame);

    KCHECK(frame.getStackRoot() == nullptr);
    KCHECK(root->getTopFrame() == nullptr);
  }

  KCHECK(melon::tryGetCurrentAsyncStackRoot() == nullptr);
}

TEST(AsyncStack, PushPop) {
  melon::detail::ScopedAsyncStackRoot scopedRoot{nullptr};

  auto& root = melon::getCurrentAsyncStackRoot();

  melon::AsyncStackFrame frame1;
  melon::AsyncStackFrame frame2;
  melon::AsyncStackFrame frame3;

  scopedRoot.activateFrame(frame1);

  KCHECK_EQ(&frame1, root.getTopFrame());
  KCHECK_EQ(&root, frame1.getStackRoot());

  melon::pushAsyncStackFrameCallerCallee(frame1, frame2);

  KCHECK_EQ(&frame2, root.getTopFrame());
  KCHECK_EQ(&frame1, frame2.getParentFrame());
  KCHECK_EQ(&root, frame2.getStackRoot());
  KCHECK(frame1.getStackRoot() == nullptr);

  melon::pushAsyncStackFrameCallerCallee(frame2, frame3);

  KCHECK_EQ(&frame3, root.getTopFrame());
  KCHECK_EQ(&frame2, frame3.getParentFrame());
  KCHECK_EQ(&frame1, frame2.getParentFrame());
  KCHECK(frame1.getParentFrame() == nullptr);
  KCHECK(frame2.getStackRoot() == nullptr);

  melon::deactivateAsyncStackFrame(frame3);

  KCHECK(root.getTopFrame() == nullptr);
  KCHECK(frame3.getStackRoot() == nullptr);

  melon::activateAsyncStackFrame(root, frame3);

  KCHECK_EQ(&frame3, root.getTopFrame());
  KCHECK_EQ(&root, frame3.getStackRoot());

  melon::popAsyncStackFrameCallee(frame3);

  KCHECK_EQ(&frame2, root.getTopFrame());
  KCHECK_EQ(&root, frame2.getStackRoot());
  KCHECK(frame3.getStackRoot() == nullptr);

  melon::popAsyncStackFrameCallee(frame2);

  KCHECK_EQ(&frame1, root.getTopFrame());
  KCHECK_EQ(&root, frame1.getStackRoot());
  KCHECK(frame2.getStackRoot() == nullptr);

  melon::deactivateAsyncStackFrame(frame1);

  KCHECK(root.getTopFrame() == nullptr);
  KCHECK(frame1.getStackRoot() == nullptr);
}

static void checkNumFrames(size_t numFrames) {
  struct Aggregator {
    std::unordered_set<melon::AsyncStackFrame*> frames;
    void operator()(melon::AsyncStackFrame* frame) { frames.insert(frame); }
  } agg;
  melon::sweepSuspendedLeafFrames(agg);

  if constexpr (!melon::kIsDebug) {
    // tracking doesn't work in non-debug-builds
    KCHECK_EQ(0, agg.frames.size());
  } else {
    KCHECK_EQ(numFrames, agg.frames.size());
  }
}

TEST(AsyncStack, suspendedLeafTracking) {
  checkNumFrames(0);

  melon::AsyncStackFrame leafFrame;
  melon::activateSuspendedLeaf(leafFrame);
  KCHECK(melon::isSuspendedLeafActive(leafFrame));
  checkNumFrames(1);

  melon::deactivateSuspendedLeaf(leafFrame);
  KCHECK(!melon::isSuspendedLeafActive(leafFrame));
  checkNumFrames(0);
}
